#ifndef PRECOND_HPP
#define PRECOND_HPP

#include <stdio.h>
#include "omp.h"
#include "my_parcsr.hpp"
#include "cuda_common.hpp"
#include "utils.hpp"

/// @brief 预条件子的基类
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据类型，可以为float或者double
/// @tparam calc_t 计算类型，可以为float或者double
/// @tparam setup_t 建立类型，可以为float或者double
template<typename idx_t, typename calc_t, typename setup_t>
class Precond
{
public:
    calc_t wgt = 1.0; // 松弛因子

    /// @brief 预条件子的构造函数
    Precond() {}

    /// @brief 预条件子的析构函数
    virtual ~Precond() {};

    /// @brief 预条件子的建立函数，需要在使用之前调用
    /// @param A 系数矩阵
    virtual void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & A) = 0;

    /// @brief 预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool zero_guess=true) = 0;

    /// @brief 设置松弛因子
    /// @param _wgt 松弛因子
    void SetWeight(const calc_t _wgt) { wgt = _wgt; }
};


/// @brief Jacobi预条件子
///        支持CPU和GPU两种实现
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class Jacobi: public Precond<idx_t, calc_t, setup_t>
{
private:
    idx_t n; // 矩阵的局部行数
    MyParCSRMatrix<idx_t, setup_t, setup_t>* A_calc; // 保存对原矩阵的引用
public:
    data_t *diag = nullptr; // 对角线元素

    /// @brief Jacobi预条件子的构造函数
    Jacobi() { this->wgt = 1.0; }

    /// @brief Jacobi预条件子的析构函数
    /// @note 释放内存
    ~Jacobi() { 
#ifdef USE_GPU
        if (diag) cuda_free(diag);
#else
        if (diag) delete [] diag;
#endif
    }

    /// @brief Jacobi预条件子的建立函数
    /// @param A 系数矩阵
    /// @note 仅仅需要提取对角线元素
    virtual void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & A) {
        this->n = A.end_row - A.beg_row;
        A_calc = &A;
#ifdef USE_GPU
        if (diag) cuda_free(diag);
        cuda_malloc((void **)(&diag), sizeof(data_t) * n);
#else
        if (diag) delete [] diag;
        diag = new data_t[n];
#endif

        // 获取 A 的对角线块
        MyCSRMatrix<idx_t, setup_t, setup_t> * A_diag = A.diag;
        idx_t * row_ptr = A_diag->row_ptr;
        idx_t * col_idx = A_diag->col_idx;
        setup_t * vals = A_diag->vals;

        // 提取对角线元素
#ifdef USE_GPU
        cuda_separate_invD(diag, row_ptr, col_idx, vals, n);
#else
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < n; i++) {
            idx_t begin = row_ptr[i], end = row_ptr[i + 1];
            for (idx_t j = begin; j < end; j++) {
                if (i == col_idx[j]) {
                    diag[i] = (data_t) (1.0 / vals[j]);
                    break;
                }
            }
        }
#endif
    }

    /// @brief Jacobi预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @param zero_guess 是否使用零初值优化
    /// @note x^{t+1} = x^{t} + wgt * D^{-1} * (b - A * x^{t})
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool zero_guess=true) {
        const idx_t nrows = b.end_row - b.beg_row;
#ifdef USE_GPU
        cuda_vec_elem_mul(x.data, b.data, diag, nrows);
#else
        if (zero_guess) {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                x.data[i] = this->wgt * b.data[i] * diag[i];
            }
        } else {
            MyParCSRVector<idx_t, calc_t> tmp(b);
            A_calc->Mult((calc_t) -1.0, x, (calc_t) 1.0, b, tmp);
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                x.data[i] += this->wgt * tmp.data[i] * diag[i];
            }
        }
#endif
    }
};

/// @brief Point Gauss-Seidel预条件子
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class PGS: public Precond<idx_t, calc_t, setup_t>
{
private:
    idx_t n = 0; // 矩阵的局部行数
    bool inited = false; // 是否已经初始化

    // 对角线元素
    data_t *val_diag = nullptr;

    // 保存对原矩阵的引用
    MyParCSRMatrix<idx_t, setup_t, setup_t> * A_calc = nullptr;

    // 向量 halo 缓冲区
    calc_t * offd_x = nullptr; 
    calc_t * x_buf = nullptr;

public:
    /// @brief PGS预条件子的构造函数
    PGS() {
        this->wgt = 1.0;
#ifdef USE_GPU
        assert(false);
#endif
    }

    /// @brief PGS预条件子的析构函数
    /// @note 释放内存
    ~PGS() {
        if (inited) {
            delete [] val_diag;
            delete [] offd_x;
            delete [] x_buf;
        }
    }

    /// @brief PGS预条件子的建立函数
    /// @param A 系数矩阵
    /// @note 提取对角线元素，初始化向量缓冲区
    virtual void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & A) {
        this->n = A.end_row - A.beg_row;
        A_calc = &A;
        if (inited) {
            delete [] val_diag;
            delete [] offd_x;
            delete [] x_buf;
        }

        // 获取 A 的对角线块
        MyCSRMatrix<idx_t, setup_t, setup_t> * A_diag = A.diag;
        idx_t * row_ptr = A_diag->row_ptr;
        idx_t * col_idx = A_diag->col_idx;
        setup_t * vals = A_diag->vals;

        // 分配内存
        val_diag = new data_t[n];
        offd_x = new calc_t[A_calc->offd->ncols];
        x_buf = new calc_t[n];

        // 提取下三角、上三角、对角线元素
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < n; i++) {
            idx_t begin = row_ptr[i], end = row_ptr[i + 1];
            for (idx_t j = begin; j < end; j++) {
                idx_t col = col_idx[j];
                setup_t v = vals[j];
                if (col == i) { // 对角线元素
                    val_diag[i] = (data_t) (1.0 / v); 
                } 
            }   
        }

        inited = true;
    }

    /// @brief PGS预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @param zero_guess 是否使用零初值优化
    /// @note 使用了零初值优化, 详见文档
    /// @note 使用了多线程优化，线程内部 Gauss-Seidel，线程间 Jacobi
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool zero_guess=true) {
        const idx_t nrows = b.end_row - b.beg_row;
        calc_t * x_data = x.data;
        const calc_t * b_data = b.data;
        const calc_t wgt = this->wgt;
        int nprocs; MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

        if (zero_guess) {
            // 前扫: D*x^{t+1} = b - L*x^{t+1}
            #pragma omp parallel
            {
                // 线程初始化，负载均衡
                int tid = omp_get_thread_num();
                int nt  = omp_get_num_threads();
                idx_t tbeg = 0, tend = 0;
                thread_load_balance(nrows, nt, tid, tbeg, tend);
                calc_t t_res; idx_t t_begin, t_end;

                // 零初值优化的前扫，只需要计算 “本线程” 对角块的部分
                for (idx_t i = tbeg; i < tend; i++) {
                    t_res = b_data[i];
                    t_begin = A_calc->diag->row_ptr[i], t_end = A_calc->diag->row_ptr[i + 1];
                    for (idx_t j = t_begin; j < t_end; j++) {
                        if (tbeg <= A_calc->diag->col_idx[j] && i > A_calc->diag->col_idx[j])
                            t_res -= ((calc_t)A_calc->diag->vals[j]) * x_data[A_calc->diag->col_idx[j]];
                    }
                    x_data[i] = wgt * t_res * val_diag[i];
                }

                // 单进程需要在此处拷贝 x_buf
                if (nprocs == 1) {
                    for (idx_t i = tbeg; i < tend; i++)
                        x_buf[i] = x_data[i];
                }
            }

            // 多进程在此处拷贝 x_buf，并更新 halo 区
            // 计算通信重叠
            if (nprocs != 1) {
                A_calc->init_update_halo(x_data, this->offd_x);
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < nrows; i++)
                    x_buf[i] = x_data[i];
                A_calc->wait_update_halo();
            }

        } else {
            // 前扫: (D + L)*x^{t+1} = b - U*x^t

            // 非零初值优化：先更新 halo 区
            if (nprocs != 1) {
                A_calc->init_update_halo(x_data, this->offd_x);
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < nrows; i++)
                    x_buf[i] = x_data[i];
                A_calc->wait_update_halo();
            }

            #pragma omp parallel
            {
                int tid = omp_get_thread_num();
                int nt  = omp_get_num_threads();
                idx_t tbeg = 0, tend = 0;
                thread_load_balance(nrows, nt, tid, tbeg, tend);
                calc_t t_res; idx_t t_begin, t_end;

                if (nprocs == 1) {
                    for (idx_t i = tbeg; i < tend; i++)
                        x_buf[i] = x_data[i];
                }
                
                // 非零初值优化的前扫
                // 线程内 Gauss-Seidel
                // 线程间 Jacobi
                // 进程间 Jacobi
                for (idx_t i = tbeg; i < tend; i++) {
                    t_res = b_data[i];
                    t_begin = A_calc->diag->row_ptr[i], t_end = A_calc->diag->row_ptr[i + 1];
                    for (idx_t j = t_begin; j < t_end; j++) {
                        if (tbeg <= A_calc->diag->col_idx[j] && A_calc->diag->col_idx[j] < tend)
                            t_res -= ((calc_t)A_calc->diag->vals[j]) * x_buf[A_calc->diag->col_idx[j]];
                        else
                            t_res -= ((calc_t)A_calc->diag->vals[j]) * x_data[A_calc->diag->col_idx[j]];
                    }
                    t_begin = A_calc->offd->row_ptr[i], t_end = A_calc->offd->row_ptr[i + 1];
                    for (idx_t j = t_begin; j < t_end; j++) {
                        t_res -= ((calc_t)A_calc->offd->vals[j]) * this->offd_x[A_calc->offd->col_idx[j]];
                    }
                    x_buf[i] += wgt * t_res * val_diag[i];
                }

                if (nprocs == 1) {
                    #pragma omp barrier
                    for (idx_t i = tbeg; i < tend; i++)
                        x_data[i] = x_buf[i];
                }
            }

            if (nprocs != 1) {
                A_calc->init_update_halo(x_buf, this->offd_x);
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < nrows; i++)
                    x_data[i] = x_buf[i];
                A_calc->wait_update_halo();
            }
        }

        // 后扫: (D+U)*x^{t+2} = b - L*x^{t+1}
        #pragma omp parallel
        {
            int tid = omp_get_thread_num();
            int nt  = omp_get_num_threads();
            idx_t tbeg = 0, tend = 0;
            thread_load_balance(nrows, nt, tid, tbeg, tend);
            calc_t t_res; idx_t t_begin, t_end;

            // 后扫，一定是非零初值
            // 线程内 Gauss-Seidel
            // 线程间 Jacobi
            // 进程间 Jacobi
            for (idx_t i = tend - 1; i >= tbeg; i--) {
                t_res = b_data[i];
                t_begin = A_calc->diag->row_ptr[i], t_end = A_calc->diag->row_ptr[i + 1];
                for (idx_t j = t_end - 1; j >= t_begin; j--) {
                    if (tbeg <= A_calc->diag->col_idx[j] && A_calc->diag->col_idx[j] < tend)
                        t_res -= ((calc_t)A_calc->diag->vals[j]) * x_data[A_calc->diag->col_idx[j]];
                    else
                        t_res -= ((calc_t)A_calc->diag->vals[j]) * x_buf[A_calc->diag->col_idx[j]];
                }
                t_begin = A_calc->offd->row_ptr[i], t_end = A_calc->offd->row_ptr[i + 1];
                for (idx_t j = t_end - 1; j >= t_begin; j--) {
                    t_res -= ((calc_t)A_calc->offd->vals[j]) * this->offd_x[A_calc->offd->col_idx[j]];
                }
                x_data[i] += wgt * t_res * val_diag[i];
            }
        }
    }
};



#endif // PRECOND_HPP
